#ifndef __UTILS_HPP__
#define __UTILS_HPP__

#include <bits/stdc++.h>
#include <unistd.h>
#include <infiniband/verbs.h>
#include "memtool.hpp"
#include <getopt.h>

#define RSEC_NETWORK_IB 1
#define RSEC_NETWORK_ROCE 2

//change if it is ROCE
#define RSEC_NETWORK_MODE RSEC_NETWORK_IB


#define ROLE_CLIENT 0
#define ROLE_SERVER 1
#define RSEC_CQ_DEPTH 128
#define MAX_IB_SGE_SIZE 4
#define MAX_INLINE 256
#define PSN 3185

#define STR_LEN 64

#define TEST_WRITE 0
#define TEST_READ 1
#define TEST_SEND 3
#define UNUSED(x) (void)(x)

typedef unsigned long long ull;

const char *ROLE_STR[2]={"client","server"};

typedef struct input_params {
  int32_t base_port_index;
  int32_t num_qps;
  int32_t device_id;
  int32_t role_int;
  char memcache_ip[STR_LEN];
  int32_t memcache_port;
  int32_t blk_size;
  char test_mode[STR_LEN];
  int64_t iter_num;
} input_params;

typedef struct ib_qp_attr{
  char name[STR_LEN];
  int32_t lid;
  int32_t qpn;
} ib_qp_attr;

typedef struct ib_mr_attr{
  uint64_t addr;
  uint32_t rkey;
} ib_mr_attr;


typedef struct ib_inf {
  input_params *inp;
  struct ibv_context *ctx;
  struct ibv_pd *pd;
  struct ibv_port_attr port_attr;
  struct ibv_device_attr dev_attr;
  //local qp and cq
  struct ibv_qp **local_qp;
  struct ibv_cq *send_cq;
  struct ibv_cq *recv_cq;
  //memcache
  MemcacheInfo *mi;
  //remote attr
  ib_qp_attr **remote_attr;
  int test_mode;
} ib_inf;



struct ib_inf *ib_complete_setup(struct input_params *conf){
  struct ib_inf *node_inf;
  node_inf = (struct ib_inf*)malloc(sizeof(ib_inf));
  assert(node_inf);
  node_inf->inp = conf;
  int res;

  node_inf->mi = mt_create_info(conf->memcache_ip, conf->memcache_port);
  assert(node_inf->mi);

  int local_rcqps;
  if(conf->role_int == ROLE_CLIENT){
    local_rcqps = conf->num_qps;
    mt_set(node_inf->mi, "num_qps", (char *)&local_rcqps, sizeof(local_rcqps));  
    mt_set(node_inf->mi, "iter", (char *)&conf->iter_num, sizeof(conf->iter_num));
  }else {
    int *tmp_qps;
    while(mt_get(node_inf->mi, "num_qps", (char **)&tmp_qps, NULL))usleep(1000);
    local_rcqps = *tmp_qps;
    conf->num_qps = local_rcqps;
    cout<<"num of qps is "<<local_rcqps<<endl;
    free(tmp_qps);

    int64_t *tmp_iter;
    while(mt_get(node_inf->mi, "iter", (char **)&tmp_iter, NULL))usleep(1000);
    conf->iter_num = *tmp_iter;
    cout<<"iteration is "<<conf->iter_num<<endl;
  }

  //get the device to use and alloc the pd
  {
    struct ibv_device **dev_list;
    int num_devices;
    dev_list = ibv_get_device_list(&num_devices);
    assert(conf->device_id < num_devices);
    node_inf->ctx = ibv_open_device(dev_list[conf->device_id]);
    assert(node_inf->ctx);
    res = ibv_query_device(node_inf->ctx, &node_inf->dev_attr);
    assert(res == 0);
    res = ibv_query_port(node_inf->ctx, conf->base_port_index, &node_inf->port_attr);
    assert(res == 0);
    node_inf->pd = ibv_alloc_pd(node_inf->ctx);
  }
  //create qps
  {
    node_inf->recv_cq = ibv_create_cq(node_inf->ctx, RSEC_CQ_DEPTH, NULL, NULL, 0);
    node_inf->send_cq = ibv_create_cq(node_inf->ctx, RSEC_CQ_DEPTH, NULL, NULL, 0);
    node_inf->local_qp = (struct ibv_qp **)malloc(sizeof(struct ibv_qp*)*local_rcqps);
    assert(node_inf->recv_cq && node_inf->send_cq && node_inf->local_qp);
    for(int i = 0;i < local_rcqps ; i++){
      struct ibv_qp_init_attr c_attr;
      memset(&c_attr, 0, sizeof(struct ibv_qp_init_attr));
      c_attr.send_cq = node_inf->send_cq;
      c_attr.recv_cq = node_inf->recv_cq;
      c_attr.qp_type = IBV_QPT_RC;

      c_attr.cap.max_send_wr = RSEC_CQ_DEPTH;
      c_attr.cap.max_recv_wr = RSEC_CQ_DEPTH;
      c_attr.cap.max_send_sge = MAX_IB_SGE_SIZE;
      c_attr.cap.max_recv_sge = MAX_IB_SGE_SIZE;
      c_attr.cap.max_inline_data = 0; //MAX_INLINE;
      c_attr.sq_sig_all = 0;

      node_inf->local_qp[i] = ibv_create_qp(node_inf->pd, &c_attr);
      assert(node_inf->local_qp[i]);

      struct ibv_qp_attr i_attr;
      memset(&i_attr,0,sizeof(struct ibv_qp_attr));
      i_attr.qp_state = IBV_QPS_INIT;
      i_attr.port_num = conf->base_port_index;
      i_attr.qp_access_flags = IBV_ACCESS_LOCAL_WRITE | 
                                IBV_ACCESS_REMOTE_WRITE |
                                IBV_ACCESS_REMOTE_READ |
                                IBV_ACCESS_REMOTE_ATOMIC;
      res = ibv_modify_qp(node_inf->local_qp[i], &i_attr ,IBV_QP_STATE | 
                          IBV_QP_PKEY_INDEX | IBV_QP_PORT |
                          IBV_QP_ACCESS_FLAGS);
      assert(res == 0);
    }
  }

  //publish all qps
  {
    for(int i = 0; i < local_rcqps ; i++){
      char name[STR_LEN];
      sprintf(name, "%s-%d", ROLE_STR[conf->role_int],i);
      ib_qp_attr qp_attr;
      strcpy(qp_attr.name, name);
      //get lid
      qp_attr.lid = node_inf->port_attr.lid;
      qp_attr.qpn = node_inf->local_qp[i]->qp_num;
      mt_set(node_inf->mi, name, (const char *)&qp_attr, sizeof(ib_qp_attr));
    }
  }
  //get remote attr
  {
    node_inf->remote_attr = (ib_qp_attr **)malloc(sizeof(ib_qp_attr *) * local_rcqps);
    for(int i = 0; i < local_rcqps ; i++ ){
      char name[STR_LEN];
      int remote_role = conf->role_int == ROLE_CLIENT ? 
                        ROLE_SERVER : ROLE_CLIENT;
      sprintf(name, "%s-%d", ROLE_STR[remote_role], i);
      while(mt_get(node_inf->mi, name, (char**)&node_inf->remote_attr[i], NULL))usleep(1000);
    }
  }
  //connect qp
  {
    for(int i = 0; i < local_rcqps ; i++){
      struct ibv_qp_attr attr;
      attr.qp_state = IBV_QPS_RTR;
      attr.path_mtu = IBV_MTU_4096;
      attr.dest_qp_num = node_inf->remote_attr[i]->qpn;
      attr.rq_psn = PSN;
      attr.max_dest_rd_atomic = 16;  // relate to IBV_WC_REM_INV_REQ_ERR
      attr.min_rnr_timer = 12;
      attr.ah_attr.is_global = 0;
      attr.ah_attr.dlid = node_inf->remote_attr[i]->lid;
      attr.ah_attr.sl = 0;
      attr.ah_attr.src_path_bits = 0;
      attr.ah_attr.port_num = conf->base_port_index;
      res = ibv_modify_qp(node_inf->local_qp[i], &attr,
              IBV_QP_STATE | IBV_QP_AV | IBV_QP_PATH_MTU |
              IBV_QP_DEST_QPN | IBV_QP_RQ_PSN |
              IBV_QP_MAX_DEST_RD_ATOMIC | IBV_QP_MIN_RNR_TIMER
              );
      assert(res == 0);

      attr.qp_state = IBV_QPS_RTS;
      attr.timeout = 14;
      attr.retry_cnt = 7;
      attr.rnr_retry = 7;
      attr.sq_psn = PSN;
      attr.max_rd_atomic = 16; 
      res = ibv_modify_qp(node_inf->local_qp[i], &attr,
              IBV_QP_STATE | IBV_QP_TIMEOUT | IBV_QP_RETRY_CNT |
               IBV_QP_RNR_RETRY | IBV_QP_SQ_PSN |
              IBV_QP_MAX_QP_RD_ATOMIC);
      assert(res == 0);
    }
  }

  return node_inf;
}


#endif